package com.ittao.taoaicodemother.service.impl;

import com.ittao.taoaicodemother.service.TaskManagerService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 任务管理服务实现类 - 改进版本，支持彻底中断
 */
@Service
@Slf4j
public class TaskManagerServiceImpl implements TaskManagerService {

    // 存储任务的中断状态
    private final ConcurrentHashMap<String, AtomicBoolean> taskInterruptFlags = new ConcurrentHashMap<>();

    // 存储任务的Disposable，用于真正取消订阅
    private final ConcurrentHashMap<String, Disposable> taskDisposables = new ConcurrentHashMap<>();

    // 存储任务的Sink，用于手动完成流
    private final ConcurrentHashMap<String, Sinks.Many<String>> taskSinks = new ConcurrentHashMap<>();

    @Override
    public Flux<String> startGenerateTask(String taskId, Long appId, Flux<String> originalFlux) {
        log.info("启动生成任务，taskId: {}, appId: {}", taskId, appId);

        // 初始化中断标志
        AtomicBoolean interruptFlag = new AtomicBoolean(false);
        taskInterruptFlags.put(taskId, interruptFlag);

        // 创建可手动控制的Sink
        Sinks.Many<String> sink = Sinks.many().multicast().onBackpressureBuffer();
        taskSinks.put(taskId, sink);

        // 订阅原始流，但要能够被中断
        Disposable disposable = originalFlux
            .takeWhile(item -> !interruptFlag.get()) // 检查中断标志
            .doOnNext(item -> {
                // 双重检查中断状态
                if (!interruptFlag.get()) {
                    Sinks.EmitResult result = sink.tryEmitNext(item);
                    if (result.isFailure()) {
                        log.warn("向sink发送数据失败，taskId: {}, result: {}", taskId, result);
                    }
                } else {
                    log.info("任务已被中断，停止发送数据，taskId: {}", taskId);
                }
            })
            .doOnComplete(() -> {
                if (!interruptFlag.get()) {
                    log.info("生成任务正常完成，taskId: {}", taskId);
                    sink.tryEmitComplete();
                } else {
                    log.info("生成任务被中断后完成，taskId: {}", taskId);
                }
                cleanupTask(taskId);
            })
            .doOnError(error -> {
                log.error("生成任务出错，taskId: {}, error: {}", taskId, error.getMessage());
                if (!interruptFlag.get()) {
                    sink.tryEmitError(error);
                }
                cleanupTask(taskId);
            })
            .doOnCancel(() -> {
                log.info("生成任务被取消，taskId: {}", taskId);
                interruptFlag.set(true); // 设置中断标志
                cleanupTask(taskId);
            })
            .subscribe();

        // 保存Disposable，用于真正的中断
        taskDisposables.put(taskId, disposable);

        // 返回sink的流，同时监听下游的取消事件
        return sink.asFlux()
            .doOnCancel(() -> {
                log.info("下游取消订阅，中断任务，taskId: {}", taskId);
                interruptTaskInternal(taskId);
            })
            .doFinally(signalType -> {
                log.info("任务流结束，signalType: {}, taskId: {}", signalType, taskId);
                // 确保清理资源
                cleanupTask(taskId);
            });
    }

    @Override
    public boolean interruptTask(String taskId, Long appId) {
        log.info("中断生成任务，taskId: {}, appId: {}", taskId, appId);
        return interruptTaskInternal(taskId);
    }

    /**
     * 内部中断任务方法
     */
    private boolean interruptTaskInternal(String taskId) {
        AtomicBoolean interruptFlag = taskInterruptFlags.get(taskId);
        Disposable disposable = taskDisposables.get(taskId);
        Sinks.Many<String> sink = taskSinks.get(taskId);

        if (interruptFlag == null) {
            log.warn("尝试中断不存在的任务，taskId: {}", taskId);
            return false;
        }

        try {
            // 1. 设置中断标志
            boolean wasInterrupted = interruptFlag.getAndSet(true);
            if (wasInterrupted) {
                log.info("任务已经被中断过了，taskId: {}", taskId);
                return true;
            }

            // 2. 取消原始流的订阅（这是关键，能真正停止AI生成）
            if (disposable != null && !disposable.isDisposed()) {
                disposable.dispose();
                log.info("成功取消原始流订阅，taskId: {}", taskId);
            }

            // 3. 完成sink流
            if (sink != null) {
                Sinks.EmitResult result = sink.tryEmitComplete();
                log.info("发送完成信号到sink，taskId: {}, result: {}", taskId, result);
            }

            // 4. 立即清理资源
            cleanupTask(taskId);

            log.info("成功中断任务，taskId: {}", taskId);
            return true;

        } catch (Exception e) {
            log.error("中断任务时发生异常，taskId: {}, error: {}", taskId, e.getMessage(), e);
            return false;
        }
    }

    @Override
    public boolean isTaskInterrupted(String taskId) {
        AtomicBoolean interruptFlag = taskInterruptFlags.getOrDefault(taskId, new AtomicBoolean(true));
        return interruptFlag.get();
    }

    @Override
    public void cleanupTask(String taskId) {
        log.info("清理任务资源，taskId: {}", taskId);

        try {
            // 清理中断标志
            taskInterruptFlags.remove(taskId);

            // 清理并释放Disposable
            Disposable disposable = taskDisposables.remove(taskId);
            if (disposable != null && !disposable.isDisposed()) {
                disposable.dispose();
                log.debug("清理时释放了Disposable，taskId: {}", taskId);
            }

            // 清理Sink
            Sinks.Many<String> sink = taskSinks.remove(taskId);
            if (sink != null) {
                // 尝试完成sink（如果还没有完成的话）
                sink.tryEmitComplete();
                log.debug("清理时完成了Sink，taskId: {}", taskId);
            }

        } catch (Exception e) {
            log.warn("清理任务资源时发生异常，taskId: {}, error: {}", taskId, e.getMessage());
        }
    }

    @Override
    public String getTaskStatus(String taskId) {
        if (taskInterruptFlags.containsKey(taskId)) {
            return isTaskInterrupted(taskId) ? "INTERRUPTED" : "RUNNING";
        }
        return "NOT_FOUND";
    }
}
